import datetime
import warnings
from typing import NewType

from dateutil import parser, rrule

__all__ = ["DATE_FORMAT",
           "m2d",
           "d2m",
           "get_date_by_month",
           "get_date_by_week",
           "get_date_by_year",
           "segment_week_by_month",
           "date_add_days",
           "get_day_interval",
           "search_trade_day",
           "interp_trading_days",
           "divide_over_10_years",
           "format_day_interval",
           "format_month_interval",
           "DateUtility",
           ]

DATE_FORMAT = "%Y-%m-%d"
MONTH_TO_DAYS = 30.42


def m2d(m):
    return int(m * MONTH_TO_DAYS)


def d2m(d):
    return int(d / MONTH_TO_DAYS)


class DateUtility:
    def __init__(self, trading_dates):
        self.trading_dates = trading_dates
        self.interp_nature_map = interp_trading_days(trading_dates)

    def get_date_by_month(self):
        pass

    def get_interval_trading_index(self, start_idx, date_interval, d_format="m"):
        # note: find_idx can be none if out of range
        if d_format == "m":
            interval = m2d(date_interval)
        elif d_format == "d":
            interval = date_interval
        add_day = date_add_days(self.trading_dates[start_idx], interval)
        find_idx = self.interp_nature_map.get(add_day)
        return find_idx

    def get_trading_day(self, index):
        return self.trading_dates[index]


def interp_trading_days(trading_dates):
    """
    interp between trading dates in nature day. so can faster indexing value of trading dates
    :param trading_dates: trading series of days: str
    :return: nature_days mapping index of trading dates.
    """
    assert len(trading_dates) > 0
    start_day = trading_dates[0]
    end_day = trading_dates[-1]
    rule_days = rrule.rrule(rrule.DAILY, dtstart=parser.parse(start_day), until=parser.parse(end_day))
    result_dict = {}
    trading_day_index = 0
    for day in rule_days:
        day_str = day.strftime(DATE_FORMAT)
        if trading_dates[trading_day_index] == day_str:
            result_dict[day_str] = trading_day_index
            trading_day_index += 1
        else:
            result_dict[day_str] = trading_day_index - 1
        if trading_day_index >= len(trading_dates):
            break
    return result_dict


def get_date_by_month(date_series: list, record_date: int, out_format="%Y-%m-%d") -> "index list and date list":
    """
    给定一个record date, 从开始日期到结束日期的每一个月，得到每一个月的这个记录天的最近有记录天(有时不开门)，以及相应索引
    :param  date_series: list of time date like 2021-09-25
    :param  record_date: day ex. 25
    :param out_format: date format
    :return: index list, record list by month
    """
    parser_time_series = list(map(parser.parse, date_series))
    assert isinstance(record_date, int), AssertionError("argument: record date should be int type")
    assert 1 <= record_date <= 28, AssertionError("date should 1<= date <=28")
    assert parser_time_series[0] - parser_time_series[1] < datetime.timedelta(days=0), \
        AssertionError("time series should ascending")
    start_date = parser_time_series[0]
    rule_date = list(rrule.rrule(rrule.MONTHLY,
                                 dtstart=start_date.replace(day=record_date), until=parser_time_series[-1]))
    if start_date > rule_date[0]:
        rule_date = rule_date[1:]
    index_result = []
    month_date_result = []
    idx_rule = 0
    for idx, iter_date in enumerate(parser_time_series):
        if iter_date - rule_date[idx_rule] >= datetime.timedelta(0):
            if iter_date - rule_date[idx_rule] == datetime.timedelta(0):
                index_result.append(idx)
                month_date_result.append(parser.parse(date_series[idx]).strftime(out_format))
            else:
                index_result.append(idx - 1)
                month_date_result.append(parser.parse(date_series[idx - 1]).strftime(out_format))
            idx_rule += 1
        if idx_rule > len(rule_date) - 1:
            break
    return index_result, month_date_result


def get_date_by_week(date_series: list, record_weekday: int) -> "index list and date list":
    """
    给定一个record date, 从开始日期到结束日期的每一个周，得到每一个周的这个记录天的记录(可能不开门就漏掉)，以及相应索引

    :param date_series: list of time date like 2021-09-25
    :param record_weekday: day: Monday == 0 ... Sunday == 6 \n
    :return: index list , record list by week \n
    """
    assert isinstance(record_weekday, int), AssertionError("argument: record date should be int type")
    index_result = []
    week_date_result = []
    for idx, date in enumerate(date_series):
        trans_date = parser.parse(date)
        if trans_date.weekday() == record_weekday:
            index_result.append(idx)
            week_date_result.append(date)
    return index_result, week_date_result


def get_date_by_year(date_series, step=1):
    now_date = parser.parse(date_series[-1])
    now_month = now_date.month
    step_day = step * 240  # one year 240 work day
    date_tail = len(date_series) - 1
    result_index = []
    result_date = []
    for idx in range(date_tail, 0, -step_day):
        date = parser.parse(date_series[idx])
        if date.month == now_month:
            res_idx = idx
        elif date.month > now_month:
            forward_idx = idx - 1
            while parser.parse(date_series[forward_idx]).month != now_month and forward_idx > 0:
                forward_idx -= 1
            if forward_idx == 0:
                continue
            res_idx = forward_idx
        else:
            backward_idx = idx + 1
            while parser.parse(date_series[backward_idx]).month != now_month and backward_idx < date_tail:
                backward_idx += 1
            res_idx = backward_idx
        result_index.append(res_idx)
        result_date.append(parser.parse(date_series[res_idx]).strftime("%Y-%m"))
    return result_index[::-1], result_date[::-1]


# TODO: should into DateSeries Class
def segment_week_by_month(month_index_list: list, week_index_list: list):
    """
    mapping week in month bucket like {17: [1, 5, 10, 15], 35: [20, 22, 27, 32]} for followp week to month

    :param month_index_list: list of month index, sort less to more
    :param week_index_list: list of week index, sort less to more
    :return: month bucket like  {17: [1, 5, 10, 15], 35: [20, 22, 27, 32]}
    """
    week_pt = 0
    month_pt = 0
    month_bucket = {}
    while week_pt < len(week_index_list) and month_pt < len(month_index_list):
        month_index = month_index_list[month_pt]
        week_index = week_index_list[week_pt]
        if week_index < month_index:  # [ ) left include right exclude.
            if month_bucket.get(month_index) is None:
                month_bucket[month_index] = [week_index]
            else:
                month_bucket[month_index].append(week_index)
            week_pt += 1
        else:
            month_pt += 1
    return month_bucket


def search_trade_day(trading_dates, nature_day) -> (int, str):
    """
    :param trading_dates: list of str "%Y-%m-%d"
    :param nature_day: nature day of str "%Y-%m-%d"
    :return: list of (found index in trading dates, str date "%Y-%m-%d")
    note: if nature day is out of range in trading dates, return index is -1
    note: the faster way is to interp nature days between trading days. so use this function will warn.
    """
    warnings.warn("function search_trade_day should be deprecated. It can be replaced by interp day.")
    assert len(trading_dates) > 0
    parser_time_series = list(map(parser.parse, trading_dates))
    parser_nature = parser.parse(nature_day)
    increasing_flag = (parser_time_series[-1] > parser_time_series[0])
    if not increasing_flag:
        parser_time_series = parser_time_series[::-1]
    if parser_nature < parser_time_series[0]:
        warnings.warn("nature_day smaller than given trading day")
        return -1, parser_time_series[0].strftime(DATE_FORMAT)
    if parser_nature > parser_time_series[-1]:
        warnings.warn("nature_day larger than given trading day")
        return -1, parser_time_series[-1].strftime(DATE_FORMAT)
    tail = len(parser_time_series) - 1
    head = 0
    while head < tail:
        mid = (head + tail) // 2 + 1  # round always forward
        if parser_nature < parser_time_series[mid]:
            tail = mid - 1
        elif parser_nature > parser_time_series[mid]:
            head = mid
        else:
            head = mid
            break
    return head, trading_dates[head]


def date_add_days(nature_day, day_count):
    added_day = parser.parse(nature_day) + datetime.timedelta(day_count)
    return added_day.strftime(DATE_FORMAT)


def get_day_interval(start_day, end_day) -> int:
    interval = parser.parse(end_day) - parser.parse(start_day)
    return interval.days


def format_day_interval(day_interval: int, o_format: str = "year") -> str:
    """
    format int day intervals to string of interval
    :param day_interval: number of days
    :param o_format year | day | month
    :return: string. if format=year return string like "约6年2月";if day return string like "2624天"
    """
    if o_format == "day":
        res_str = f"{day_interval}天"
    elif o_format == "year":
        year = day_interval // 365
        month = int((day_interval % 365) / MONTH_TO_DAYS)
        res_str = f"{year}年{month}月"
    elif o_format == "month":
        month = int(day_interval / MONTH_TO_DAYS)
        res_str = f"{month}月"
    else:
        raise RuntimeError("error format in get_date_interval function: year | day | month")
    return res_str


def format_month_interval(month_interval: int, o_format: str = "year") -> str:
    """
    format int day intervals to string of interval
    :param month_interval: number of month
    :param o_format year | day | month
    :return: string. if format=year return string like "约6年2月";if day return string like "2624天"
    """
    if o_format == "day":
        res_str = f"{m2d(month_interval)}天"
    elif o_format == "year":
        year = month_interval // 12
        month = month_interval % 12
        res_str = f"{year}y"
    elif o_format == "month":
        res_str = f"{month_interval}月"
    else:
        raise RuntimeError("error format in get_date_interval function: year | day | month")
    return res_str


def divide_over_10_years(time_start: str, time_end: str) -> list:
    """
    divide over 10 years time range to some time segments of interval-10y.
    :param time_start: string
    :param time_end: string
    :return: return a list of tuple[start, end], each of which has 10-year-interval.
    """
    date_start = parser.parse(time_start)
    date_end = parser.parse(time_end)
    date_range = list(rrule.rrule(10*rrule.YEARLY, dtstart=date_start, until=date_end))[::10]
    result = []
    for i in range(len(date_range)-1):
        result.append((date_range[i], date_range[i+1]-datetime.timedelta(1)))
    if date_end != date_range[-1]:
        result.append((date_range[-1], date_end))
    return list(map(lambda x: tuple(x_ele.strftime(DATE_FORMAT) for x_ele in x), result))


if __name__ == "__main__":
    from investment.time_series import IndexTimeSeries

    my_index = IndexTimeSeries('399995', "2020-03-31", "2021-03-31", fill_flag=True)
    dates = my_index.dates_series
    month_indexes, month_date = get_date_by_month(dates, 25)
    test_month = ['2020-04-24', '2020-05-25', '2020-06-24', '2020-07-24', '2020-08-25', '2020-09-25',
                  '2020-10-23', '2020-11-25', '2020-12-25', '2021-01-25', '2021-02-25', '2021-03-25']
    assert (month_date == test_month)
    week_indexes, week_date = get_date_by_week(dates, 2)  # 2 is Wed.
    bucket = segment_week_by_month(month_indexes, week_indexes)
    print(week_date)
